import json
import random

from agents.base_agent import BaseAgent
from prompts.paper_writer_prompts import NEW_CHAPTER_WRITING_SYSTEM_PROMPT, NEW_CHAPTER_WRITING_USER_PROMPT
from my_utils.outline_parser import OutlineUtils


class PaperWriter(BaseAgent):
    def __init__(self, max_epochs: int = 20):
        super().__init__(max_epochs=max_epochs)


async def write_paper(outline_dict: dict, paper_index_dict: dict, selected_chunks: dict) -> tuple[dict, list]:
    """ 根据大纲内容，生成综述论文的每个小节内容 """

    utils = OutlineUtils()
    # 获取所有叶子节点（最底层节点）信息
    leaf_nodes = utils.get_all_leaf_nodes(outline_dict)

    # 将简洁ID替换为实际的论文ID
    for item in leaf_nodes:
        if "refs" in item:
            item["refs"] = [paper_index_dict.get(ref) for ref in item["refs"] if paper_index_dict.get(ref) is not None]
        else:
            item["refs"] = []
    # 调试输出所有叶子节点信息
    print("获取所有叶子节点信息：")
    for node in leaf_nodes:
        print(f"标题: {node['title']}")
        print(f"内容: {node['content']}")
        print(f"引用: {node['refs']}\n")

    # 创建引用ID列表，用于后续引用编号
    ref_id_list = []
    for item in leaf_nodes:
        if len(item["refs"]) > 10:
            item["refs"] = random.sample(item["refs"], 10)
        for idx, ref_id in enumerate(item["refs"]):
            # 组合论文ID和chunk_id
            ref_id_with_chunk_id = f"{ref_id}-{selected_chunks[ref_id]['entity']['chunk_id']}"
            if ref_id_with_chunk_id not in ref_id_list:
                ref_id_list.append(ref_id_with_chunk_id)

            # 获取引用的序号（从1开始）并构建引用块结构
            ref_index = ref_id_list.index(ref_id_with_chunk_id) + 1
            item["refs"][idx] = (f'<paper_chunk {idx + 1}>\n'
                                 f'<paper_id>paper_{ref_index}</paper_id>\n'
                                 f'<paper_content>{selected_chunks[ref_id]["entity"]["chunk_text"]}</paper_content>\n'
                                 f'</paper_chunk {idx + 1}>\n')

    # 保存引用ID列表到文件
    with open("examples_data/ref_id_list.json", "w", encoding='utf-8') as f:
        json.dump(ref_id_list, f, ensure_ascii=False, indent=4)

    # 获取大纲目录结构
    user_prompt_template = NEW_CHAPTER_WRITING_USER_PROMPT
    outline = utils.get_table_of_contents(outline_dict)

    # 为每个小节生成提示内容
    user_input_list = []
    for item in leaf_nodes:
        ref_list = "\n".join([f'",\n{ref}' for ref in item["refs"]])
        section_prompt = user_prompt_template.format(
            section_title=item["title"],
            outline=outline,
            ref_list=ref_list
        )
        user_input_list.append(section_prompt)

    # 调用大模型生成各节内容
    paper_writer_agent = PaperWriter()
    result = await paper_writer_agent.run_loop(
        model_name="GLM_4_PLUS",
        system_prompt=NEW_CHAPTER_WRITING_SYSTEM_PROMPT,
        user_inputs=user_input_list,
        temperature=0.1,
        max_wait_times=1000,
        max_tokens=20480
    )

    # 更新叶子节点的内容
    for i, item in enumerate(leaf_nodes):
        item["content"] = result[i].replace("\\\\n", "\n")

    # 更新整个大纲结构
    updated_outline_recursive_dict = utils.update_leaf_nodes(outline_dict, leaf_nodes)

    # 自底向上填充父节点内容
    filled_outline_recursive_dict = utils.fill_parent_contents(updated_outline_recursive_dict)

    return filled_outline_recursive_dict, ref_id_list

if __name__ == "__main__":
    # 读取大纲xml文件和paper_index_dict
    with open("../examples_data/outline_xml.txt", "r", encoding='utf-8') as f:
        outline_xml = f.read()
    with open("../examples_data/paper_index_dict.json", "r", encoding='utf-8') as f:
        paper_index_dict = json.load(f)

    with open('examples_data/origin_paper_search_result.json', 'r', encoding='utf-8') as f:
        selected_chunks = json.load(f)

    filled_outline_recursive_dict = write_paper(outline_xml, paper_index_dict, selected_chunks)

    # 保存filled_outline_recursive_dict到文件中
    with open("../outputs/filled_outline_recursive_dict.json", "w", encoding='utf-8') as f:
        json.dump(filled_outline_recursive_dict, f, ensure_ascii=False, indent=4)

    print("1")
